{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -*- coding: utf-8 -*-\n",
    "from miasm.analysis.sandbox import Sandbox_Linux_x86_64\n",
    "from miasm.analysis.binary import Container\n",
    "from miasm.os_dep.linux_stdlib import linobjs\n",
    "from miasm.core.utils import hexdump\n",
    "from miasm.core.locationdb import LocationDB\n",
    "from miasm.analysis.dse import DSEPathConstraint\n",
    "from miasm.analysis.machine import Machine\n",
    "from miasm.expression.expression import ExprMem, ExprId, ExprInt, ExprAssign\n",
    "from future.utils import viewitems\n",
    "from miasm.jitter.csts import PAGE_READ, PAGE_WRITE\n",
    "import sys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = Sandbox_Linux_x86_64.parser(description='ELF sandboxer')\n",
    "options = parser.parse_args(args=[])\n",
    "options.filename = 'flattening_volatile.bin'\n",
    "options.strategy = 'code-cov'\n",
    "options.mimic_env = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Need to adjust the context of dse to the one of concrete\n",
    "def xxx_printf_symb(dse):\n",
    "\n",
    "    regs = dse.ir_arch.arch.regs\n",
    "    ret_addr = ExprInt(dse.jitter.get_stack_arg(0), regs.RIP.size)\n",
    "\n",
    "    dse.update_state({\n",
    "        regs.RSP: dse.symb.eval_expr(regs.RSP + ExprInt(8, regs.RSP.size)),\n",
    "        dse.ir_arch.IRDst: ret_addr,\n",
    "        regs.RAX: ExprInt(6, regs.RAX.size),\n",
    "        regs.RIP: ret_addr,\n",
    "    })"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def code_sentinelle(jitter):\n",
    "    jitter.run = False\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_db = LocationDB()\n",
    "sb = Sandbox_Linux_x86_64(loc_db, options.filename, options, globals())\n",
    "\n",
    "with open(options.filename, 'rb') as fstream:\n",
    "    cont = Container.from_stream(fstream, loc_db)\n",
    "    \n",
    "machine = Machine('x86_64')\n",
    "\n",
    "ret_addr = 0x000000001337beef\n",
    "sb.jitter.add_breakpoint(ret_addr, code_sentinelle)\n",
    "sb.jitter.push_uint64_t(ret_addr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# init_run need to be placed before DSE is attacked\n",
    "sb.jitter.init_run(0x1040)\n",
    "\n",
    "MEM_ARGV_ADDR  = 0x7ff70000\n",
    "MEM_ARGV1_ADDR = 0x7ff80000\n",
    "\n",
    "# argv\n",
    "sb.jitter.vm.add_memory_page(\n",
    "        MEM_ARGV_ADDR,\n",
    "        PAGE_READ | PAGE_WRITE,\n",
    "        b'\\x42\\x42\\x42\\x42\\x42\\x42\\x42\\x42' +  \n",
    "        b'\\x00\\x00\\xf8\\x7f\\x00\\x00\\x00\\x00', \n",
    "        'Binary'\n",
    "    )\n",
    "\n",
    "# argv[1]\n",
    "sb.jitter.vm.add_memory_page(\n",
    "        MEM_ARGV1_ADDR,\n",
    "        PAGE_READ | PAGE_WRITE,\n",
    "        b'\\x00\\x00\\x00\\x00\\x00\\x00\\x00\\x00',\n",
    "        'Binary'\n",
    "    )\n",
    "\n",
    "sb.jitter.cpu.RDI = 0x2             # argc\n",
    "sb.jitter.cpu.RSI = MEM_ARGV_ADDR   # argv\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert strategy to the correct value\n",
    "strategy = {\n",
    "    'code-cov': DSEPathConstraint.PRODUCE_SOLUTION_CODE_COV,\n",
    "    'branch-cov': DSEPathConstraint.PRODUCE_SOLUTION_BRANCH_COV,\n",
    "    'path-cov': DSEPathConstraint.PRODUCE_SOLUTION_PATH_COV,\n",
    "}[options.strategy]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Init a DSE instance with a given strategy\n",
    "dse = DSEPathConstraint(machine, loc_db, produce_solution=strategy)\n",
    "dse.attach(sb.jitter)\n",
    "dse.update_state_from_concrete()\n",
    "dse.add_lib_handler(sb.libs, globals())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Symbolize the argument\n",
    "regs = sb.jitter.ir_arch.arch.regs\n",
    "\n",
    "argv1 = [] \n",
    "for i in range(8):\n",
    "    # Create an ExprId for argv[1][x]\n",
    "    argv1.append(ExprId('Argv[1][%d]'%(i), 8))\n",
    "\n",
    "    # Add constraints because argv1[x] is a readable ascii\n",
    "    const = dse.z3_trans.from_expr(argv1[i])\n",
    "    dse.cur_solver.add(31 < const) \n",
    "    dse.cur_solver.add(const < 127)\n",
    "\n",
    "argv1_addr = []\n",
    "for i in range(8):\n",
    "    argv1_addr.append(ExprMem(ExprInt(MEM_ARGV1_ADDR + i, 64), 8))\n",
    "\n",
    "s = {}\n",
    "for addr, argv in zip(argv1_addr, argv1):\n",
    "    s[addr] = argv\n",
    "    \n",
    "dse.update_state(s)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "todo = set([(ExprInt(0x42, 8), \n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8),\n",
    "            ExprInt(0x41, 8))])\n",
    "done = set()\n",
    "\n",
    "snapshot = dse.take_snapshot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "while todo:\n",
    "    arg_value = todo.pop()\n",
    "\n",
    "    if arg_value in done:\n",
    "        continue\n",
    "\n",
    "    done.add(arg_value)\n",
    "\n",
    "    for i in range(8):\n",
    "        print('Run with ARG = %s' % (arg_value[i]))\n",
    "\n",
    "    print('---start---')\n",
    "\n",
    "    dse.restore_snapshot(snapshot, keep_known_solutions=True)\n",
    "\n",
    "    for i in range(8):\n",
    "        sb.jitter.eval_expr(ExprAssign(argv1_addr[i], arg_value[i]))\n",
    "\n",
    "    sb.jitter.init_run(0x1040)\n",
    "    sb.jitter.set_trace_log(trace_regs=True, trace_new_blocks=False)\n",
    "    sb.jitter.continue_run(step=False)\n",
    "\n",
    "    print('---end---')\n",
    "\n",
    "    if sb.jitter.cpu.RAX == 0:\n",
    "        print('FOUND!!!')\n",
    "        argv1_str = ''.join([chr(x) for x in arg_value])\n",
    "        print(argv1_str)\n",
    "        break\n",
    "\n",
    "    for sol_ident, model in viewitems(dse.new_solutions):\n",
    "\n",
    "        print('Found a solution to reach: %s' % str(sol_ident))\n",
    "        print('model', model)\n",
    "\n",
    "        sol_expr = []\n",
    "        for i in range(8):\n",
    "            # Get the argument to use as a Miasm Expr\n",
    "            try:\n",
    "                sol_value = model.eval(dse.z3_trans.from_expr(argv1[i])).as_long()\n",
    "            except AttributeError:\n",
    "                sol_value = 0\n",
    "\n",
    "            sol_expr.append(ExprInt(sol_value, argv1[i].size))\n",
    "\n",
    "            print('\\tARG[1][%d] = %s' % (i, sol_expr[i]))\n",
    "\n",
    "        todo.add((sol_expr[0], \n",
    "                    sol_expr[1], \n",
    "                    sol_expr[2], \n",
    "                    sol_expr[3], \n",
    "                    sol_expr[4],\n",
    "                    sol_expr[5], \n",
    "                    sol_expr[6], \n",
    "                    sol_expr[7]))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
